// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once

#include <memory>
#include <string>
#include <unordered_map>
#include <utility>

#include <boost/date_time/gregorian/gregorian.hpp>
#include <fmt/format.h>

// internal
#include "coin/feed/fastfeed/feed.pb.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/market/market_exchange_api.h"
#include "orio/io/base.h"
#include "orio/io/concat.h"
#include "orio/io/file.h"
#include "orio/record/simple.h"

namespace coin2::feed::fastfeed {
namespace impl {

using ::coin2::exchange::base::market::MarketExchangeApi;
using ::fastfeed::proto::Feed;

class FastFeedWriter {
 public:
  FastFeedWriter(
      const std::string& root_dir,  // may or may not end with '/'
      const std::string& machine,
      const MarketExchangeApi& mea,
      const std::string& recipe,
      int worker,
      const std::string& extension,
      const std::string& start)
      : root_dir_(root_dir),
        machine_(machine),
        mea_(mea),
        recipe_(recipe),
        worker_(worker),
        extension_(extension) {
    CHECK_THROW(extension.empty() || extension == ".gz" || extension == ".lz4")
        << "[FlowWriter] Unsupported compression type: " + extension;
    auto date = PtimeFromIso8601(start).date();
    auto start_time = PtimeFromIso8601(start).time_of_day();

    if (start_time.hours() != 0 ||
        start_time.minutes() != 0 ||
        start_time.seconds() != 0) {
      start_time_ = fmt::format(
          ".{:s}T{:02d}{:02d}{:02d}",
          boost::gregorian::to_iso_string(date),
          start_time.hours(),
          start_time.minutes(),
          start_time.seconds());
    }
  }

  virtual ~FastFeedWriter() {
    for (auto& [key, writer] : writers_) {
      (void)key;  // unused
      writer->Flush();
    }
  }

  virtual void Write(Feed* in) {
    auto& symbol = in->symbol();
    int64_t date = in->timestamp() / (24LL * 60 * 60 * 1'000'000'000);
    orio::io::Writer<orio::record::SimpleRecord>* writer = nullptr;
    if (writers_.find(symbol) == writers_.end() || last_dates_.at(symbol) != date) {
      // {config.fastfeed_root_dir}/{machine}/date/{mea}/
      // '%s.W%s.%s' % (str(product), worker extension)

      last_dates_[symbol] = date;
      auto ptime = PtimeFromTimestamp(in->timestamp());
      auto YYYYMMDD = boost::gregorian::to_iso_string(ptime.date());
      std::experimental::filesystem::path path = fmt::format(
          "{root_dir}/{machine}/{YYYYMMDD}/{mea}.{recipe}/"
          "{symbol}{start_time_tag}.W{worker}{dot_extension}",
          fmt::arg("root_dir", root_dir_),
          fmt::arg("machine", machine_),
          fmt::arg("YYYYMMDD", YYYYMMDD),
          fmt::arg("start_time", start_time_),
          fmt::arg("mea", mea_.String()),
          fmt::arg("recipe", recipe_),
          fmt::arg("symbol", symbol),
          fmt::arg("start_time_tag", start_time_),
          fmt::arg("worker", worker_),
          fmt::arg("dot_extension", extension_));

      std::string tmp_name = fmt::format(
          "coin2_fastfeed/{machine}/{YYYYMMDD}/{mea}.{recipe}/"
          "{symbol}{start_time_tag}.W{worker}{dot_extension}",
          fmt::arg("machine", machine_),
          fmt::arg("YYYYMMDD", YYYYMMDD),
          fmt::arg("start_time", start_time_),
          fmt::arg("mea", mea_.String()),
          fmt::arg("recipe", recipe_),
          fmt::arg("symbol", symbol),
          fmt::arg("start_time_tag", start_time_),
          fmt::arg("worker", worker_),
          fmt::arg("dot_extension", extension_));
      auto tmp_path = orio::io::GetTemporaryFile(tmp_name);

      writers_[symbol] = std::unique_ptr<orio::io::Writer<orio::record::SimpleRecord>>(
          orio::record::SimpleWriter::ToFile(tmp_path, path));
      writer = writers_[symbol].get();
      LOG(INFO) << "[FastFeedWriter] creating a fastfeed file\n"
                << "dst file: " << path.c_str() << "\ntmp file: " << tmp_path; 
    } else {
      writer = writers_[symbol].get();
      // Symbol is only populated for the first update, to make the file size smaller.
      in->clear_symbol();
    }

    in->SerializeToString(&buf1_);
    buf2_.resize(buf1_.size() + sizeof(uint32_t));
    auto* record = reinterpret_cast<orio::record::SimpleRecord*>(buf2_.data());
    record->len = buf1_.size();
    std::memcpy(&record->data[0], buf1_.data(), record->len);
    writer->Write(*record);
  }

 private:
  std::unordered_map<std::string, std::unique_ptr<orio::io::Writer<orio::record::SimpleRecord>>>
      writers_;

 protected:
  std::string root_dir_;
  std::string machine_; 
  const MarketExchangeApi mea_;
  const std::string recipe_;
  int worker_;
  std::string extension_;
  std::string start_time_;
  std::unordered_map<std::string, int64_t> last_dates_;  // symbol:(epoch//1day_ns)
  std::string buf1_;
  std::string buf2_;
};

}  // namespace impl

using impl::FastFeedWriter;

}  // namespace coin2::feed::fastfeed
